// SPDX-License-Identifier: GPL-2.0

#define pr_fmt(fmt) "DIM Calculate: " fmt

#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/err.h>
#include <linux/elf.h>
#include <linux/fs.h>
#include <crypto/hash.h>
#include <linux/kallsyms.h>
#include <linux/highmem.h>
#include <linux/vmalloc.h>
#include <linux/jump_label.h>
#include <linux/sort.h>

#include "dim_elf.h"
#include "dim_calculate.h"
#include "dim_utils.h"

#define TRAMPOLINE_SECTION_NAME ".vos_patch_trampoline_seg"
#define MAX_TEXT_NUMBER 64

/* elf information for user task digest calculating */
struct elf_info {
	unsigned long text_vaddr[MAX_TEXT_NUMBER];
	unsigned long text_vsize[MAX_TEXT_NUMBER];
	unsigned long text_num;
	unsigned long trampo_vaddr;
};

static int parse_elf(struct file *elf_file,
		     struct elf_info *info,
		     int trampo_flag)
{
	int ret;
	struct elfhdr ehdr;
	struct elf_shdr shdr = {0};
	struct elf_phdr *phdrs = NULL;
	struct elf_phdr *phdr = NULL;
	size_t i, j, phdr_num;
	unsigned long temp;

	if (elf_file == NULL || info == NULL)
		return -EINVAL;

	/* 1. parse elf header */
	ret = dim_elf_get_ehdr(elf_file, &ehdr);
	if (ret < 0)
		return ret;

	/* 2. parse elf program headers */
	ret = dim_elf_get_phdrs(elf_file, &ehdr, &phdrs, &phdr_num);
	if (ret < 0)
		return ret;

	/* 3. find text segment from program headers */
	info->text_num = 0;
	for (i = 0, phdr = phdrs; i < phdr_num; i++, phdr++) {
		if ((phdr->p_type == PT_LOAD) && (phdr->p_flags & PF_R) &&
		    (phdr->p_flags & PF_X) && !(phdr->p_flags & PF_W)) {
			info->text_vaddr[info->text_num] = phdr->p_vaddr;
			info->text_vsize[info->text_num] = phdr->p_memsz;

			if (++info->text_num >= MAX_TEXT_NUMBER) {
				kfree(phdrs);
				return -ENOEXEC;
			}
		}
	}

	if (info->text_num == 0) {
		kfree(phdrs);
		return -ENOEXEC;
	}

	/* sort text segment by vaddr */
	for (i = 0; i < info->text_num - 1; i++) {
		for (j = 0; j < info->text_num - 1 - i; j++) {
			if (info->text_vaddr[j] > info->text_vaddr[j + 1]) {
				temp = info->text_vaddr[j];
				info->text_vaddr[j] = info->text_vaddr[j + 1];
				info->text_vaddr[j + 1] = temp;
				temp = info->text_vsize[j];
				info->text_vsize[j] = info->text_vsize[j + 1];
				info->text_vsize[j + 1] = temp;
			}
		}
	}

	kfree(phdrs);
	if (!trampo_flag) /* needn't to parse trampoline section */
		return info->text_num ? 0 : -ENOEXEC;

	/* 4. parse trampoline section */
	ret = elf_get_section_by_name(elf_file, &ehdr,
		TRAMPOLINE_SECTION_NAME, &shdr);
	if (ret < 0)
		info->trampo_vaddr = 0;
	else
		info->trampo_vaddr = shdr.sh_addr;

	return 0;
}

static int dim_update_vm_digest(struct mm_struct *mm,
				unsigned long start,
				unsigned long len,
				struct shash_desc *shash)
{
	long page_num_get;
	size_t page_num, page_offset_last, i;
	struct page **pages = NULL;
	void *page_ptr = NULL;

	page_num = (len + PAGE_SIZE - 1) / PAGE_SIZE;
	page_offset_last = len % PAGE_SIZE;
	page_offset_last = page_offset_last == 0 ?
		PAGE_SIZE : page_offset_last;

	pages = vzalloc(page_num * sizeof(struct page *));
	if (pages == NULL)
		return -ENOMEM;

	page_num_get = get_user_pages_remote(mm, start, page_num, 0,
		pages, NULL, NULL);
	if (page_num_get < 0) {
		pr_err("get phy pages failed: %ld\n", page_num_get);
		vfree(pages);
		return page_num_get;
	}

	for (i = 0; i < page_num_get; i++) {
		page_ptr = kmap(pages[i]);
		if (page_ptr == NULL)
			break;

		crypto_shash_update(shash, page_ptr,
			i == (page_num - 1) ? page_offset_last : PAGE_SIZE);
		kunmap(pages[i]);
	}

	while (page_num_get)
		put_page(pages[--page_num_get]);
	vfree(pages);
	return 0;
}

static int dim_calc_vm_digest(struct mm_struct *mm,
			      unsigned long start,
			      unsigned long len,
			      enum dim_hash_type hash_index,
			      char *digest_out)
{
	int ret;
	struct crypto_shash *tfm = dim_hash_get_tfm(hash_index);

	DIM_SHASH_DESC_ON_STACK(shash);

	if (tfm == NULL)
		return -EOPNOTSUPP;

	shash->tfm = tfm;

	ret = crypto_shash_init(shash);
	if (ret < 0) {
		pr_warn("shash init failed: %d\n", ret);
		return ret;
	}

	dim_update_vm_digest(mm, start, len, shash);
	return crypto_shash_final(shash, digest_out);
}

static int calc_buf_digest(char *inbuf,
			   size_t len,
			   enum dim_hash_type hash_index,
			   char *digest_out)
{
	int ret;
	struct crypto_shash *tfm = dim_hash_get_tfm(hash_index);

	DIM_SHASH_DESC_ON_STACK(shash);

	if (tfm == NULL)
		return -EOPNOTSUPP;

	shash->tfm = tfm;

	ret = crypto_shash_init(shash);
	if (ret < 0) {
		pr_warn("shash init failed: %d\n", ret);
		return ret;
	}

	crypto_shash_update(shash, inbuf, len);

	ret = crypto_shash_final(shash, digest_out);
	if (ret < 0) {
		pr_warn("shash final failed: %d\n", ret);
		return ret;
	}

	return dim_hash_get_digest_len(hash_index);
}

int dim_calc_user_task_trampoline_digest(struct mm_struct *mm,
					 struct vm_area_struct *vma,
					 enum dim_hash_type hash_index,
					 char *digest_out)
{
	int ret;
	struct elf_info info;
	struct vm_area_struct *vma_trampo = NULL;
	unsigned long addr_trampo;

	if (mm == NULL || vma == NULL || dim_get_vma_file(vma) == NULL ||
	    digest_out == NULL || hash_index >= DIM_HASH_LAST)
		return -EINVAL;

	/* 1. parse elf info from elf file */
	ret = parse_elf(dim_get_vma_file(vma), &info, 1);
	if (ret < 0) {
		pr_err("fail to parse elf file from vma: %d", ret);
		return ret;
	}

	if (!info.trampo_vaddr) {
		pr_err("fail to get trampoline section from elf file");
		return -EINVAL;
	}

	/* 2. find trampoline vma */
	if (vma->vm_start < info.text_vaddr[0])
		return -EFAULT;

	addr_trampo = vma->vm_start - info.text_vaddr[0] + info.trampo_vaddr;
	for (vma_trampo = vma;
	     vma_trampo != NULL;
	     vma_trampo = vma_trampo->vm_next) {
		if (vma_trampo->vm_start > addr_trampo) {
			vma_trampo = NULL;
			break; /* trampoline section not found */
		}

		if (vma_trampo->vm_start == addr_trampo)
			break; /* trampoline vma found */
	}

	if (vma_trampo == NULL) {
		pr_err("fail to find trampoline vma");
		return -ENODATA;
	}

	/* 3. calculate digest of trampoline vma */
	ret = dim_calc_vm_digest(mm, vma_trampo->vm_start,
		vma_trampo->vm_end - vma_trampo->vm_start,
		hash_index, digest_out);
	if (ret < 0) {
		pr_err("fail to calculate digest of trampoline vma: %d", ret);
		return ret;
	}

	return dim_hash_get_digest_len(hash_index);
}

static bool is_rx_only_vma(struct vm_area_struct *vma)
{
	if ((vma->vm_flags & VM_READ) &&
	    (vma->vm_flags & VM_EXEC) &&
	    !(vma->vm_flags & VM_WRITE))
		return true;

	return false;
}

int dim_calc_user_task_text_digest(struct mm_struct *mm,
				   struct vm_area_struct *vma,
				   enum dim_hash_type hash_index,
				   char *digest_out)
{
	int ret, cal_err;
	struct elf_info info;
	unsigned long addr_end, addr_trampo;
	unsigned long vma_cal_size;
	struct vm_area_struct *vma_cur = vma;
	unsigned long text_idx = 0;
	struct crypto_shash *tfm = dim_hash_get_tfm(hash_index);

	DIM_SHASH_DESC_ON_STACK(shash);

	if (mm == NULL || vma == NULL || dim_get_vma_file(vma) == NULL ||
	    digest_out == NULL || hash_index >= DIM_HASH_LAST)
		return -EINVAL;

	if (tfm == NULL)
		return -EOPNOTSUPP;

	shash->tfm = tfm;

	ret = crypto_shash_init(shash);
	if (ret < 0) {
		pr_warn("shash init failed: %d\n", ret);
		return ret;
	}

	/* 1. parse elf info from elf file */
	ret = parse_elf(dim_get_vma_file(vma), &info, 1);
	if (ret < 0) {
		pr_err("fail to parse elf file from vma: %d", ret);
		return ret;
	}

	/* 2. find and calculate text segment digest */
	cal_err = 0;
	addr_trampo = (info.trampo_vaddr == 0) ? 0 :
		vma->vm_start - info.text_vaddr[0] + info.trampo_vaddr;
	while (vma_cur != NULL) {
		/* check text segment size right */
		addr_end = vma_cur->vm_start + info.text_vsize[text_idx];
		for (; vma_cur != NULL && vma_cur->vm_start < addr_end;
		       vma_cur = vma_cur->vm_next) {
			/* check vma attribute */
			if (!is_rx_only_vma(vma_cur) ||
			dim_get_vma_file(vma_cur) != dim_get_vma_file(vma)) {
				cal_err = 1;
				break;
			}

			vma_cal_size = vma_cur->vm_end > addr_end ?
				addr_end - vma_cur->vm_start :
				vma_cur->vm_end - vma_cur->vm_start;
			dim_update_vm_digest(mm,
				vma_cur->vm_start, vma_cal_size, shash);
		}

		if (cal_err)
			break;

		for (; vma_cur != NULL; vma_cur = vma_cur->vm_next) {
			/* skip the mmaped trampoline */
			if (vma_cur->vm_start == addr_trampo)
				continue;

			if (is_rx_only_vma(vma_cur) && dim_get_vma_file(vma_cur) != NULL) {
				if (dim_get_vma_file(vma_cur) != dim_get_vma_file(vma))
					cal_err = 1;

				break;
			}
		}

		if (++text_idx >= info.text_num || cal_err)
			break;
	}

	ret = crypto_shash_final(shash, digest_out);
	if (ret < 0) {
		pr_warn("shash final failed: %d\n", ret);
		return ret;
	}

	return dim_hash_get_digest_len(hash_index);
}

static int jcode_cmp(const void *a, const void *b)
{
	return *(unsigned long *)a - *(unsigned long *)b;
}

static int sort_jump_table_code(struct jump_entry *sjump,
				unsigned int jump_counts,
				unsigned long **code)
{
	unsigned int i;
	unsigned long *buf = NULL;

	buf = vzalloc(sizeof(unsigned long) * jump_counts);
	if (buf == NULL)
		return -ENOMEM;

	jump_label_lock();
	for (i = 0; i < jump_counts; i++)
		buf[i] = jump_entry_code(&sjump[i]);
	jump_label_unlock();

	sort(buf, jump_counts, sizeof(unsigned long), jcode_cmp, NULL);
	*code = buf;
	return 0;
}

static int calc_kernel_digest(char *saddr,
			      char *eaddr,
			      enum dim_hash_type hash_index,
			      char *digest_out)
{
	int ret;
	struct jump_entry *sjump = NULL;
	struct jump_entry *ejump = NULL;
	unsigned int jump_counts, i;
	unsigned long *code_sort = NULL;
	unsigned long jump_code;
	char *cur_addr = saddr;
	struct crypto_shash *tfm = dim_hash_get_tfm(hash_index);

	DIM_SHASH_DESC_ON_STACK(shash);

	if (tfm == NULL)
		return -EOPNOTSUPP;

	shash->tfm = tfm;

	sjump = (struct jump_entry *)kallsyms_lookup_name("__start___jump_table");
	ejump = (struct jump_entry *)kallsyms_lookup_name("__stop___jump_table");

	if (sjump == NULL || ejump == NULL || sjump >= ejump) {
		jump_counts = 0;
	} else {
		jump_counts = ((uintptr_t)ejump - (uintptr_t)sjump) /
			sizeof(struct jump_entry);
		if (jump_counts > 0) {
			ret = sort_jump_table_code(sjump, jump_counts, &code_sort);
			if (ret < 0)
				return ret;
		}
	}

	ret = crypto_shash_init(shash);
	if (ret < 0) {
		pr_warn("shash init failed, %d\n", ret);
		vfree(code_sort);
		return ret;
	}

	for (i = 0; i < jump_counts && (uintptr_t)cur_addr < (uintptr_t)eaddr; i++) {
		jump_code = code_sort[i];
		if (jump_code < (uintptr_t)cur_addr) /* jump_code can be 0 */
			continue;

		if (jump_code >= (uintptr_t)eaddr) /* no more valid jump code */
			break;

		if (jump_code > (uintptr_t)cur_addr)
			crypto_shash_update(shash, cur_addr,
				jump_code - (uintptr_t)cur_addr);

		cur_addr = (char *)(jump_code + JUMP_LABEL_NOP_SIZE);
	}

	if ((uintptr_t)cur_addr < (uintptr_t)eaddr)
		crypto_shash_update(shash, cur_addr,
			(uintptr_t)eaddr - (uintptr_t)cur_addr);

	vfree(code_sort);
	ret = crypto_shash_final(shash, digest_out);
	if (ret < 0) {
		pr_warn("shash final failed, %d\n", ret);
		return ret;
	}

	return dim_hash_get_digest_len(hash_index);
}

int dim_calc_kernel_digest(enum dim_hash_type hash_index, char *digest_out)
{
	char *stext = NULL;
	char *etext = NULL;
	int size;

	if (digest_out == NULL || hash_index >= DIM_HASH_LAST)
		return -EINVAL;

	stext = (char *)kallsyms_lookup_name("_stext");
	etext = (char *)kallsyms_lookup_name("_etext");
	if (stext == NULL || etext == NULL || stext > etext)
		return -EFAULT;

	size = calc_kernel_digest(stext, etext, hash_index, digest_out);
	if (size < 0) {
		pr_warn("calculate digest failed\n");
		return size;
	}

	return size;
}

int dim_calc_ko_digest(struct module *mod,
		       enum dim_hash_type hash_index,
		       char *digest_out)
{
	if (mod == NULL || digest_out == NULL ||
	    hash_index >= DIM_HASH_LAST)
		return -EINVAL;

	return calc_buf_digest(mod->core_layout.base,
			       mod->core_layout.text_size,
			       hash_index,
			       digest_out);
}
